#pragma once

#include "ast_config.hpp"
#include "my_parser.hpp"
#include "../../jsonxx/fifo_map.hpp"
#include <exception>
#include <map>
#include <memory>
#include <set>
#include <stack>
#include <unordered_set>
#include <vector>
#include <deque>
#include <ctime>

using nlohmann::fifo_map;
using std::deque;
using std::make_shared;
using std::map;
using std::shared_ptr;
using std::stack;
using std::unordered_map;
using std::unordered_set;
using std::vector;

class ASTNode;

using PtrASTNode = shared_ptr<ASTNode>;


class ASTNode {
public:
    string name;    // Name of the node
    fifo_map<string, vector<Token>> fields; // Fields in the node
    vector<PtrASTNode> children;    // children
    ASTNode() = default;
    void printAST(std::ostream& oss){
        ASTNode::printAST(oss, this);
    }
    /**
     * Build AST from parse tree
     * @param p parse tree root
     * @param config AST config
     */ 
    static PtrASTNode buildAST(MyParser::PtrParseTreeNode p, const ASTConfig& config) {
        cout << "[Info] Start building AST\n";
        PtrASTNode ast = buildASTStep1_rawAST(p, config);
        // cout << "[Info] Finish step 1\n";
        ASTNode::buildASTStep2_child2Field(ast, config);
        // cout << "[Info] Finish step 2\n";
        ASTNode::buildASTStep3_child2Sibling(ast, config);
        // cout << "[Info] Finish step 3\n";
        ASTNode::buildASTStep4_mergeChilds(ast, config);
        clock_t ed = clock();
        // cout << "[Info] Finish step 4\n";
        cout << "[Info] AST built successfully!\n";
        return ast;
    }
private:
    template<class Ptr>
    static void printAST(std::ostream &oss, const Ptr &t,
                         const string &indent = "") {
        oss << indent << "|-- " << t->name << "{ ";
        int cnt = 0;
        for (auto i : t->fields) {
            oss << i.first << ": ";
            if (i.second.empty())
                oss << i.second[0].word;
            else {
                for (auto j : i.second) {
                    oss << j.word << " ";
                }
            }
            ++cnt;
            if (cnt < t->fields.size())
                oss << ", ";
        }
        oss << " }\n";
        for (auto i : t->children) {
            printAST(oss, i, indent + "  ");
        }
    }

private:
    static PtrASTNode buildASTStep1_rawAST(MyParser::PtrParseTreeNode p,
                                           const ASTConfig &config) { // 第一步，建立AST
        PtrASTNode ast = make_shared<ASTNode>();
        vector<PtrASTNode> newChildren;
        ast->name = p->ruleT.getName();
        for (auto i : p->children) {
            if (!i->isTerminal()) { // 递归
                newChildren.push_back(buildASTStep1_rawAST(i, config));
            } else if (!config.ignoredTerminals.count(i->token.tokenName)) {
                // 把不可删除的终结符提升到树根
                if(!config.terminalNameMap.count(i->token.tokenName)){
                    cout << "[Warn] No terminal name mapping for " << i->token.tokenName
                         << "\n";
                    continue;
                }
                if (ast->fields.count(config.terminalNameMap.at(i->token.tokenName))) {
                    ast->fields.at(config.terminalNameMap.at(i->token.tokenName))
                        .push_back(i->token);
                } else {
                    ast->fields[config.terminalNameMap.at(i->token.tokenName)] = {
                        i->token};
                }
            }
        }
        if (config.generalization.count(p->ruleT.getName())) {
            if (newChildren.size() == 1 && !config.compressExclusion.count(p->ruleT.getName())) {
                return newChildren[0];
            }
            ast->name = config.generalization.at(p->ruleT.getName());
        }
        ast->children = newChildren;
        return ast;
    }

    static void
    buildASTStep2_child2Field(PtrASTNode ast,
                              const ASTConfig &config) { // 第二步，child to fields
        vector<PtrASTNode> newChildren;
        for (auto i : ast->children) {
            if (config.childToFields.count(i->name)) { // 压缩
                if(!config.childToFields.at(i->name).empty()) {
                    ast->fields[config.childToFields.at(i->name)] = {{i->name, i->name}};
                }
                ast->fields.insert(i->fields.begin(), i->fields.end());
                buildASTStep2_child2Field(i, config);
                newChildren.insert(newChildren.end(), i->children.begin(),
                                   i->children.end());
            } else {
                newChildren.push_back(i);
                buildASTStep2_child2Field(i, config);
            }
        }
        ast->children = newChildren;
    }

    static vector<PtrASTNode> buildASTStep3_child2Sibling(PtrASTNode ast, const ASTConfig& config){ // 第三步，child to sibling
        vector<PtrASTNode> other2bElevate; // 其它类型递归节点的养父
        vector<PtrASTNode> self2bElevate;  // 自己类型递归节点的养父
        vector<PtrASTNode> newChildren;
        unordered_set<string> recursiveChildNames;
        bool hasRecursiveChild = false;
        for (auto i:ast->children) {
            auto&& tmp = buildASTStep3_child2Sibling(i, config);    // 获得孩子的养父
            if (config.childToSibling.count(i->name)) { // 是要被提升的
                hasRecursiveChild = true;
                recursiveChildNames.insert(i->name);
                if (ast->name != i->name) {
                    other2bElevate.insert(other2bElevate.end(), tmp.begin(), tmp.end());
                } else {
                    self2bElevate.insert(self2bElevate.end(), tmp.begin(), tmp.end());
                }
            } else {
                newChildren.push_back(i);
            }
        }
        newChildren.insert(newChildren.end(), other2bElevate.rbegin(), other2bElevate.rend());
        ast->children = newChildren;
        if (config.childToSibling.count(ast->name)) { // 自己就是在递归
            self2bElevate.push_back(ast);
            return self2bElevate;
        }
        return {};
    }

    static void buildASTStep4_mergeChilds(PtrASTNode ast, const ASTConfig &config) {
        unordered_map<string, PtrASTNode> stepFathers; // <节点类型，节点的孩子>
        vector<PtrASTNode> newChildren;
        for(auto i: ast->children){
            buildASTStep4_mergeChilds(i, config);
            if (config.mergeChilds.count(i->name)) {
                if(!stepFathers.count(i->name)){
                    stepFathers[i->name] = make_shared<ASTNode>();
                }
                stepFathers[i->name]->children.insert(stepFathers[i->name]->children.end(),
                                                      i->children.begin(), i->children.end());
            } else {
                newChildren.push_back(i);
            }
        }
        for(auto [name, father]: stepFathers){
            father->name = name;
            newChildren.push_back(father);
        }
        ast->children = newChildren;
    }
};
